#! /usr/local/bin/python
import random
import math
import sys
import os
import copy

folder = "system_" + sys.argv[1]

position=500
speed=0.001
p_speed=0.5
p_mass=[0.05/(10.0**x) for x in range(0,9)]


class Body(object):
	def __init__ (self,body_data):
		self.name=body_data[0]
		self.mass=float(body_data[1])
		self.position=Position(body_data[2:5])
		self.velocity=Velocity(body_data[5:8])
		self.acceleration=Acceleration()
		

	def __del__(self):
		del self.position
		del self.velocity
		
	def create_text(self):
		self.text=open(self.name+".txt","w+")

class Position(object):
	def __init__ (self,position_data):
		self.x=float(position_data[0])
		self.y=float(position_data[1])
		self.z=float(position_data[2])

class Velocity(object):
	def __init__ (self,velocity_data):
		self.x=float(velocity_data[0])
		self.y=float(velocity_data[1])
		self.z=float(velocity_data[2])	

class Acceleration(object):
	def __init__(self):
		self.x=0
		self.y=0
		self.z=0
		
	def reset(self):
		self.x=0
		self.y=0
		self.z=0


class System(object):
	def __init__(self):
		self.bodies=[]
		self.build()
		self.stability = self.evaluate(self.bodies)
		self.printed=False
		self.evolved=0
	
	def build(self):
		bodies=self.bodies
		
		N=random.randint(2,10)
		for i in range(0,N):
			body_data=[]
			body_data.append("body_" + str(i))
			body_data.append(random.uniform(.1,15))
			for j in range(0,3):
				body_data.append(random.uniform(-position,position))
			for j in range(0,3):
				body_data.append(random.uniform(-speed,speed))
			
			body=Body(body_data)
			bodies.append(body)
	
		planet_n=random.randint(20,40)
		for i in range(0,planet_n):
			body_data=[]
			body_data.append("planet_" + str(i))
			body_data.append(random.sample(p_mass,1)[0])
			for j in range(0,3):
				body_data.append(random.uniform(-position,position))
			for j in range(0,3):
				body_data.append(random.uniform(-p_speed,p_speed))
			
			body=Body(body_data)
			bodies.append(body)
		

	def mutate(self, alphaMass, alphaPosition, alphaVelocity):
		oldSelf = copy.deepcopy(self)
		sample = random.sample(range(0,len(self.bodies)-1),5)
		for whichBody in sample:
			

			self.bodies[whichBody].mass += alphaMass * random.uniform(-5,5)
			if self.bodies[whichBody].mass <= 0 or self.bodies[whichBody].mass > 5:
				self.bodies[whichBody].mass = oldSelf.bodies[whichBody].mass
			
			self.bodies[whichBody].velocity.x+= alphaVelocity * (random.uniform(-speed,speed))
			self.bodies[whichBody].velocity.y+= alphaVelocity * (random.uniform(-speed,speed))
			self.bodies[whichBody].velocity.z+= alphaVelocity * (random.uniform(-speed,speed))
			


			self.bodies[whichBody].position.x+= alphaPosition * (random.uniform(-position,position))
			if self.bodies[whichBody].position.x < -position or self.bodies[whichBody].position.x > position:
				self.bodies[whichBody].position.x = oldSelf.bodies[whichBody].position.x
			
			self.bodies[whichBody].position.y+= alphaPosition * (random.uniform(-position,position))
			if self.bodies[whichBody].position.y < -position or self.bodies[whichBody].position.y > position:
				self.bodies[whichBody].position.y = oldSelf.bodies[whichBody].position.y

			self.bodies[whichBody].position.z+= alphaPosition * (random.uniform(-position,position))
			if self.bodies[whichBody].position.z < -position or self.bodies[whichBody].position.z > position:
				self.bodies[whichBody].position.z = oldSelf.bodies[whichBody].position.z
		
		self.stability = self.evaluate(self.bodies)
		


		if self.stability > oldSelf.stability :
			self = copy.deepcopy(oldSelf)

							     
	def evaluate(self,bodies):
		kinetic=0.0
		potential=0.0
		
		G=2.93558*10**-4
		
		for body in bodies:
			vel = body.velocity
			vel_sq = (vel.x**2 + vel.y**2 + vel.z**2)
			kinetic += 0.5*body.mass*vel_sq
		
		for i in range(0,len(bodies)):
			current_body=bodies[i]
			current_position=current_body.position
		
			for j in range(0,i):
				other_body=bodies[j]
				other_position=other_body.position		


				d_x=(other_position.x-current_position.x)
				d_y=(other_position.y-current_position.y)
				d_z=(other_position.z-current_position.z)
			
				radius = (d_x**2 + d_y**2 + d_z**2)**(0.5)
				
				potential -= G*current_body.mass*other_body.mass/radius
		try:
			return round(kinetic/potential*(-2.0))
		except:
			return 100.0


class GeneticAlgorithm:
	def __init__(self,populationSize,maxGeneration):
		systems=[]
		self.maxGen = maxGeneration
		for i in range(0,populationSize):
			systems.append(System())
		systems = self.selection(systems)
		self.genetitize(populationSize,systems)

	def mutateAll(self):
		for system in systems:
			system.mutate(system,.1,.1,.1)
	
	def selection(self,systems):
		val_checker=[]

		for system in systems:
			val_checker.append(system.stability)
		
		for system in systems:
			if system.stability == 1:
				Evolve(system,0,True)
	
		val_checker.sort()
		print val_checker
		
		systems=systems[0:4]
		
		for system in systems:
			m=random.uniform(-1,1)
			system.mutate(m*0.001,m*10,m*10)
				
		return systems
		
	def genetitize(self,populationSize,systems):
		for i in range(0,self.maxGen):
			for i in range(len(systems),populationSize):
				systems.append(System())
			self.selection(systems)
			
	def write_conditions(self,system):
		os.chdir("data")
		os.mkdir(folder)
		os.chdir(folder)
		for body in system.bodies:
			pos=body.position
			vel=body.velocity
			data_file = open(body.name,"w+")
			data_file.write(str(body.mass)+"\t"+str(pos.x)+"\t"+str(pos.y)+"\t"+str(pos.z)+"\t"+str(vel.x)+"\t"+str(vel.y)+"\t"+str(vel.z) +"\n")



class Evolve:
	def __init__(self,system):
		if keep_system:
			os.mkdir("evolved")
			os.chdir("evolved")
			for body in system.bodies:
				body.create_text()
		
		self.control(system)
		
	def control(self,system,t_max=100000):
		dt=0.2
		t=0
		count=1
		self.accelerate(system)
		
		for body in system.bodies:
			self.calculate_velocity(body,dt/2)
			self.calculate_position(body,dt)
			t+=dt
		
		body.acceleration.reset()
		
		while t<t_max:
			self.accelerate(system)

			for body in system.bodies:
				self.calculate_velocity(body,dt)
				self.calculate_position(body,dt)
			body.acceleration.reset()
			
			if count==1000:
				self.write_data(system)
				count=0
			t+=dt
			count+=1
			
				
	
	
	def accelerate(self,system):
		bodies = system.bodies
		G=2.93558*10**-4
		for i in range(0,len(bodies)):
			current_body=bodies[i]
			current_position=current_body.position
			for j in range(0,i):
				other_body=bodies[j]
				other_position=other_body.position
			
				d_x=(other_position.x-current_position.x)
				d_y=(other_position.y-current_position.y)
				d_z=(other_position.z-current_position.z)
			
				radius = d_x**2 + d_y**2 + d_z**2
				
				grav_mag = G/(radius**(3.0/2.0))
				
				grav_x=grav_mag*d_x
				grav_y=grav_mag*d_y
				grav_z=grav_mag*d_z
				
				current_body.acceleration.x += grav_x*other_body.mass
				other_body.acceleration.x -= grav_x*current_body.mass
				
				current_body.acceleration.y += grav_y*other_body.mass
				other_body.acceleration.y -= grav_y*current_body.mass
				
				current_body.acceleration.z += grav_z*other_body.mass
				other_body.acceleration.z -= grav_z*current_body.mass

			


	def calculate_velocity(self,body,dt):
		body.velocity.x += dt*body.acceleration.x
		body.velocity.y += dt*body.acceleration.y
		body.velocity.z += dt*body.acceleration.z

	def calculate_position(self,body,dt):
		body.position.x += dt*body.velocity.x
		body.position.y += dt*body.velocity.y
		body.position.z += dt*body.velocity.z
	
	def write_data(self,system):
		for body in system.bodies:
			pos=body.position
			vel=body.velocity
			body.text.write(str(pos.x)+"\t"+str(pos.y)+"\t"+str(pos.z)+"\n")
	
GeneticAlgorithm(10,5000000)
